import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

from loader import load_flow

class flowDataset(Dataset):
    def __init__(self, dataset_fname, load_time=0, load_length=1, load_direction=1, no_neg=True, truncate_num=100):
        self.trainset = []
        self.load_time = load_time
        self.load_length = load_length
        self.load_direction = load_direction
        self.no_neg = no_neg
        self.truncate_num = truncate_num
        self.max_label = 0
        with open(dataset_fname, 'r') as f:
            for line in f:
                site, fname_list = line.split(": ")
                site = int(site[4:]) - 1
                if site > self.max_label:
                    self.max_label = site
                fname_list = fname_list.strip()[2:-2].split("', '")
                for fname in fname_list:
                    self.trainset.append((site, fname))
        # np.random.shuffle(self.trainset)
    
    def __getitem__(self, index):
        site, fname = self.trainset[index]
        flow = load_flow(fname, load_time=self.load_time, load_length=self.load_length, load_direction=self.load_direction, no_neg=self.no_neg)
        flow = np.array(flow)
        if len(flow) > self.truncate_num:
            flow = flow[:self.truncate_num]
        elif len(flow) < self.truncate_num:
            flow = np.pad(flow, ((0, self.truncate_num - flow.shape[0])), 'constant')
        flow = flow.astype(np.int64)
        flow = torch.from_numpy(flow)
        return site, flow, fname
    
    # def __iter__(self):
    #     for site, fname in self.trainset:
    #         flow = load_flow(fname, load_time=self.load_time, load_length=self.load_length, load_direction=self.load_direction)
    #         flow = np.array(flow)
    #         if len(flow) > self.truncate_num:
    #             flow = flow[:self.truncate_num]
    #         elif len(flow) < self.truncate_num:
    #             flow = np.pad(flow, ((0, self.truncate_num - flow.shape[0])), 'constant')
    #         flow = flow.astype(np.float32)
    #         flow = torch.from_numpy(flow)
    #         yield (site, flow, fname)
    
    def __len__(self):
        return len(self.trainset)

